"""
三班排班处理模块 - 处理generate_and_save_schedule接口后的排班结果
功能：将护理员排班分为三班（每班8小时），处理跨天情况，实现多周期排班，并预留资源
使用表：schedule_results, time_slots, caregiver_schedules
返回：优化后的排班结果保存到数据库
"""
import random
import uuid
from datetime import datetime, date, timedelta
import pandas as pd
import logging
import json
from typing import List, Dict, Tuple, Set
from .models import db, ScheduleResult, TimeSlot, CaregiverSchedule, Attendantskills, ShiftAssignment
from sqlalchemy import func
from config_params import params
from contextlib import contextmanager

# 设置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class ShiftScheduler:
    """
    三班排班处理类 - 实现三班排班处理的核心逻辑
    功能：处理排班结果，分配班次，处理跨天，实现资源预留
    """
    def __init__(self, config_params=None):
        """
        初始化方法 - 设置排班参数
        参数：config_params(配置参数字典，如果为None则使用全局params)
        """
        self.params = config_params if config_params else params
        self.shift_start_time = self._parse_time(self.params['shift_start_time'])
        self.shifts_per_day = self.params['shifts_per_day']
        self.hours_per_shift = self.params['hours_per_shift']
        self.day_night_weight = self.params['day_night_weight']
        self.emergency_reserve_ratio = self.params['emergency_reserve_ratio']
        self.schedule_period_type = self.params['schedule_period_type']
        self.schedule_period_length = self.params['schedule_period_length']
        
        # 定义班次时间范围
        self.shifts = self._generate_shifts()
        
        # 初始化护理人员数据
        self.caregivers = self._load_caregivers()
        
        # 维护护理人员工作量统计
        self.caregiver_workload = {caregiver.id: 0 for caregiver in self.caregivers}
        
        # 跟踪每个班次的护理人员分配情况
        self.shift_assignments = {}
        
    def _parse_time(self, time_str: str) -> Tuple[int, int]:
        """
        解析时间字符串为小时和分钟
        参数：time_str(时间字符串，格式为HH:MM)
        返回：(小时, 分钟)元组
        """
        hours, minutes = map(int, time_str.split(':'))
        return hours, minutes
    
    def _generate_shifts(self) -> List[Dict]:
        """
        生成班次时间范围
        返回：班次列表，每个班次包含开始和结束时间
        """
        shifts = []
        start_hour, start_minute = self.shift_start_time
        
        for i in range(self.shifts_per_day):
            # 计算班次开始时间
            shift_start = start_hour + i * self.hours_per_shift
            shift_start_hour = shift_start % 24  # 处理跨天情况
            is_next_day = shift_start >= 24
            
            # 计算班次结束时间
            shift_end = shift_start + self.hours_per_shift
            shift_end_hour = shift_end % 24  # 处理跨天情况
            is_end_next_day = shift_end >= 24
            
            # 将时间转换为时间槽索引
            start_slot = int((shift_start_hour * 60 + start_minute) / 30)
            end_slot = int((shift_end_hour * 60 + start_minute) / 30)
            
            # 创建班次信息
            shift = {
                'name': f"Shift {i+1}",
                'start_hour': shift_start_hour,
                'end_hour': shift_end_hour,
                'start_slot': start_slot,
                'end_slot': end_slot,
                'is_start_next_day': is_next_day,
                'is_end_next_day': is_end_next_day,
                'is_night': shift_start_hour >= 18 or shift_start_hour < 6  # 晚上6点到早上6点为夜班
            }
            shifts.append(shift)
        
        return shifts
    
    def _load_caregivers(self) -> List:
        """
        从数据库加载护理人员数据
        返回：护理人员列表
        """
        return Attendantskills.query.all()
    
    def _get_period_dates(self, start_date: date) -> List[date]:
        """
        根据排班周期类型和长度获取所有需要排班的日期
        参数：start_date(开始日期)
        返回：需要排班的日期列表
        """
        dates = []
        current_date = start_date
        
        # 根据周期类型确定天数
        if self.schedule_period_type == 'day':
            days = self.schedule_period_length
        elif self.schedule_period_type == 'week':
            days = self.schedule_period_length * 7
        elif self.schedule_period_type == 'month':
            days = self.schedule_period_length * 30
        elif self.schedule_period_type == 'quarter':
            days = self.schedule_period_length * 90
        elif self.schedule_period_type == 'year':
            days = self.schedule_period_length * 365
        else:
            days = 1  # 默认为1天
        
        # 生成日期列表
        for i in range(days):
            dates.append(current_date)
            current_date += timedelta(days=1)
        
        return dates
    
    def _assign_caregivers_to_shifts(self, schedule_date: date) -> Dict:
        """
        为指定日期的班次分配护理人员
        参数：schedule_date(排班日期)
        返回：班次分配结果字典
        """
        # 获取当天所有时间槽的排班
        time_slots = TimeSlot.query.join(ScheduleResult).filter(
            func.date(ScheduleResult.schedule_date) == schedule_date
        ).all()
        
        # 按时间槽分组
        slots_by_time = {}
        for slot in time_slots:
            time_key = slot.time_slot
            if time_key not in slots_by_time:
                slots_by_time[time_key] = []
            slots_by_time[time_key].append(slot)
        
        # 初始化班次分配结果
        shift_assignments = {shift['name']: {'caregivers': set(), 'slots': []} for shift in self.shifts}
        
        # 遍历每个时间槽，分配到对应班次
        for time_key, slots in slots_by_time.items():
            # 解析时间槽格式，例如"08:00-08:30"
            start_time, _ = time_key.split('-')
            hours, minutes = map(int, start_time.split(':'))
            slot_index = (hours * 60 + minutes) // 30
            
            # 找出这个时间槽所属的班次
            assigned_shift = None
            for shift in self.shifts:
                # 考虑跨天情况
                if shift['is_end_next_day'] and shift['start_slot'] > shift['end_slot']:
                    # 跨天班次
                    if slot_index >= shift['start_slot'] or slot_index < shift['end_slot']:
                        assigned_shift = shift
                        break
                else:
                    # 不跨天班次
                    if shift['start_slot'] <= slot_index < shift['end_slot']:
                        assigned_shift = shift
                        break
            
            if assigned_shift:
                shift_name = assigned_shift['name']
                # 记录该班次的时间槽
                shift_assignments[shift_name]['slots'].append(time_key)
                
                # 为该班次记录护理人员
                for slot in slots:
                    if slot.caregiver_id != '-1':  # 排除未分配护理人员的时间槽
                        shift_assignments[shift_name]['caregivers'].add(slot.caregiver_id)
        
        return shift_assignments
    
    def _balance_workload(self, shift_assignments: Dict) -> Dict:
        """
        平衡各班次的工作量
        参数：shift_assignments(班次分配字典)
        返回：平衡后的班次分配字典
        """
        # 计算各班次的护理人员数量
        shift_caregiver_counts = {name: len(data['caregivers']) for name, data in shift_assignments.items()}
        
        # 计算平均每班应有的护理人员数量
        total_caregivers = sum(shift_caregiver_counts.values())
        avg_caregivers_per_shift = total_caregivers / len(self.shifts)
        
        # 尝试平衡各班次的护理人员数量
        # 从人员多的班次移动到人员少的班次
        shifts_sorted = sorted(shift_assignments.keys(), key=lambda x: len(shift_assignments[x]['caregivers']), reverse=True)
        
        for i in range(len(shifts_sorted)-1):
            source_shift = shifts_sorted[i]
            target_shift = shifts_sorted[-1]
            
            # 如果源班次和目标班次的人员差距超过阈值
            if len(shift_assignments[source_shift]['caregivers']) - len(shift_assignments[target_shift]['caregivers']) > 1:
                # 从源班次中选择一个负载较轻的护理人员
                if shift_assignments[source_shift]['caregivers']:
                    caregiver_to_move = min(shift_assignments[source_shift]['caregivers'], 
                                          key=lambda cid: self.caregiver_workload.get(cid, 0))
                    
                    # 移动护理人员
                    shift_assignments[source_shift]['caregivers'].remove(caregiver_to_move)
                    shift_assignments[target_shift]['caregivers'].add(caregiver_to_move)
                    
                    # 更新排序
                    shifts_sorted = sorted(shift_assignments.keys(), key=lambda x: len(shift_assignments[x]['caregivers']), reverse=True)
        
        return shift_assignments
    
    def _apply_day_night_weight(self, shift_assignments: Dict) -> Dict:
        """
        应用白天/黑夜权重比
        参数：shift_assignments(班次分配字典)
        返回：应用权重后的班次分配字典
        """
        # 将班次分为白班和夜班
        day_shifts = [shift['name'] for shift in self.shifts if not shift['is_night']]
        night_shifts = [shift['name'] for shift in self.shifts if shift['is_night']]
        
        # 计算当前白班和夜班的护理人员总数
        day_caregivers = set()
        night_caregivers = set()
        
        for shift_name in day_shifts:
            day_caregivers.update(shift_assignments[shift_name]['caregivers'])
        
        for shift_name in night_shifts:
            night_caregivers.update(shift_assignments[shift_name]['caregivers'])
        
        # 计算理想的白班/夜班人员比例
        total_caregivers = len(day_caregivers) + len(night_caregivers)
        ideal_day_count = int(total_caregivers * self.day_night_weight)
        ideal_night_count = total_caregivers - ideal_day_count
        
        # 如果当前分配与理想分配相差太大，进行调整
        if abs(len(day_caregivers) - ideal_day_count) > 1:
            if len(day_caregivers) > ideal_day_count:
                # 从白班移动到夜班
                caregivers_to_move = list(day_caregivers)[:len(day_caregivers) - ideal_day_count]
                
                # 移动护理人员
                for caregiver_id in caregivers_to_move:
                    # 从白班中移除
                    for shift_name in day_shifts:
                        if caregiver_id in shift_assignments[shift_name]['caregivers']:
                            shift_assignments[shift_name]['caregivers'].remove(caregiver_id)
                            break
                    
                    # 添加到夜班
                    if night_shifts:  # 确保有夜班
                        # 选择护理人员最少的夜班
                        target_shift = min(night_shifts, key=lambda x: len(shift_assignments[x]['caregivers']))
                        shift_assignments[target_shift]['caregivers'].add(caregiver_id)
            else:
                # 从夜班移动到白班
                caregivers_to_move = list(night_caregivers)[:len(night_caregivers) - ideal_night_count]
                
                # 移动护理人员
                for caregiver_id in caregivers_to_move:
                    # 从夜班中移除
                    for shift_name in night_shifts:
                        if caregiver_id in shift_assignments[shift_name]['caregivers']:
                            shift_assignments[shift_name]['caregivers'].remove(caregiver_id)
                            break
                    
                    # 添加到白班
                    if day_shifts:  # 确保有白班
                        # 选择护理人员最少的白班
                        target_shift = min(day_shifts, key=lambda x: len(shift_assignments[x]['caregivers']))
                        shift_assignments[target_shift]['caregivers'].add(caregiver_id)
        
        return shift_assignments
    
    def _reserve_emergency_resource(self, shift_assignments: Dict) -> Dict:
        """
        预留应急资源
        参数：shift_assignments(班次分配字典)
        返回：预留资源后的班次分配字典
        """
        # 计算每个班次应预留的护理人员数量
        for shift_name, data in shift_assignments.items():
            caregiver_count = len(data['caregivers'])
            reserve_count = max(1, int(caregiver_count * self.emergency_reserve_ratio))
            
            # 标记预留的护理人员
            data['reserved_caregivers'] = set()
            
            # 选择工作量较轻的护理人员作为预留
            if data['caregivers']:
                caregivers_sorted = sorted(list(data['caregivers']), key=lambda cid: self.caregiver_workload.get(cid, 0))
                data['reserved_caregivers'] = set(caregivers_sorted[:reserve_count])
        
        return shift_assignments
    
    def _update_time_slots(self, schedule_date: date, shift_assignments: Dict) -> None:
        """
        根据班次分配更新数据库中的班次分配表
        参数：schedule_date(排班日期), shift_assignments(班次分配字典)
        """
        # 获取当天所有时间槽
        time_slots = TimeSlot.query.join(ScheduleResult).filter(
            func.date(ScheduleResult.schedule_date) == schedule_date
        ).all()
        
        # 创建时间槽到班次的映射
        slot_to_shift = {}
        for shift_name, data in shift_assignments.items():
            for time_key in data['slots']:
                slot_to_shift[time_key] = {
                    'name': shift_name,
                    'is_night': any(shift['name'] == shift_name and shift['is_night'] for shift in self.shifts),
                    'start_hour': next((shift['start_hour'] for shift in self.shifts if shift['name'] == shift_name), 0),
                    'end_hour': next((shift['end_hour'] for shift in self.shifts if shift['name'] == shift_name), 0)
                }
        
        # 清除当天的旧数据
        try:
            ShiftAssignment.query.filter(
                func.date(ShiftAssignment.schedule_date) == schedule_date
            ).delete(synchronize_session=False)
            db.session.commit()
            logger.info(f"已清除日期 {schedule_date} 的旧班次分配记录")
        except Exception as e:
            db.session.rollback()
            logger.error(f"删除旧的班次分配记录失败: {str(e)}")
        
        # 创建新的班次分配记录
        new_assignments = []
        for slot in time_slots:
            shift_info = slot_to_shift.get(slot.time_slot)
            if shift_info:
                # 检查该护理员是否被预留
                is_reserved = False
                if shift_info['name'] in shift_assignments:
                    shift_data = shift_assignments[shift_info['name']]
                    if 'reserved_caregivers' in shift_data and slot.caregiver_id in shift_data['reserved_caregivers']:
                        is_reserved = True
                
                # 创建新记录
                new_assignment = ShiftAssignment(
                    id=str(uuid.uuid4()),
                    time_slot_id=slot.id,
                    elder_id=slot.elder,
                    caregiver_id=slot.caregiver_id,
                    shift_name=shift_info['name'],
                    shift_start_hour=shift_info['start_hour'],
                    shift_end_hour=shift_info['end_hour'],
                    is_night_shift=shift_info['is_night'],
                    is_reserved=is_reserved,
                    schedule_date=slot.schedule_date
                )
                new_assignments.append(new_assignment)
        
        # 批量保存记录
        if new_assignments:
            try:
                db.session.bulk_save_objects(new_assignments)
                db.session.commit()
                logger.info(f"成功添加 {len(new_assignments)} 条班次分配记录")
            except Exception as e:
                db.session.rollback()
                logger.error(f"保存班次分配记录失败: {str(e)}")
        else:
            logger.warning(f"没有找到需要保存的班次分配记录")
    
    @staticmethod
    @contextmanager
    def safe_transaction(session):
        """
        安全的数据库事务上下文管理器
        参数：session(数据库会话)
        """
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            logger.error(f"数据库事务失败: {str(e)}")
            raise e
    
    def _apply_csp_constraints(self, shift_assignments: Dict, schedule_date: date, day_index: int) -> Dict:
        """
        应用约束满足问题(CSP)算法处理护理员工作-休息限制
        参数：
            shift_assignments: 班次分配字典
            schedule_date: 排班日期
            day_index: 在排班周期中的天数索引
        返回：
            应用约束后的班次分配字典
        """
        # 如果是多天排班的第一天，则跳过约束检查
        if day_index == 0:
            return shift_assignments
            
        # 获取前一天的排班情况
        previous_date = schedule_date - timedelta(days=1)
        previous_assignments = {}
        
        # 查询前一天的ShiftAssignment记录
        previous_records = ShiftAssignment.query.filter(
            func.date(ShiftAssignment.schedule_date) == previous_date
        ).all()
        
        # 整理前一天每个护理员的班次分配
        caregiver_prev_shifts = {}
        for record in previous_records:
            caregiver_prev_shifts[record.caregiver_id] = {
                'shift_name': record.shift_name,
                'is_night_shift': record.is_night_shift
            }
        
        # 将当前班次分配复制一份用于调整
        adjusted_assignments = {shift: data.copy() for shift, data in shift_assignments.items()}
        
        # 应用硬约束：晚班后不应立即安排早班
        for shift_name, data in shift_assignments.items():
            caregivers_to_reassign = set()
            
            # 检查每个护理员的前一天班次
            for caregiver_id in data['caregivers']:
                prev_shift = caregiver_prev_shifts.get(caregiver_id, {})
                
                # 如果前一天是晚班，而今天是早班，则需要重新分配
                if prev_shift.get('is_night_shift', False) and shift_name == 'Shift 1':
                    caregivers_to_reassign.add(caregiver_id)
                    logger.info(f"需要重新分配: 护理员 {caregiver_id} 前一天是晚班，今天不应安排早班")
            
            # 从当前班次中移除需要重新分配的护理员
            if caregivers_to_reassign:
                adjusted_assignments[shift_name]['caregivers'] = data['caregivers'] - caregivers_to_reassign
                
                # 尝试将这些护理员分配到其他班次
                available_shifts = [s for s in self.shifts if s['name'] != 'Shift 1']
                if available_shifts:
                    for caregiver_id in caregivers_to_reassign:
                        # 优先分配到中班
                        target_shift = 'Shift 2' if 'Shift 2' in adjusted_assignments else 'Shift 3'
                        adjusted_assignments[target_shift]['caregivers'].add(caregiver_id)
                        logger.info(f"重新分配: 护理员 {caregiver_id} 从早班调整到 {target_shift}")
        
        # 应用软约束：尽量固定护理员在相似的班次
        caregivers_by_pref = {}
        for shift_name, data in adjusted_assignments.items():
            for caregiver_id in data['caregivers']:
                if caregiver_id not in caregivers_by_pref:
                    caregivers_by_pref[caregiver_id] = []
                caregivers_by_pref[caregiver_id].append(shift_name)
        
        # 对于分配到多个班次的护理员，尝试保持在之前的班次
        for caregiver_id, shifts in caregivers_by_pref.items():
            if len(shifts) > 1 and caregiver_id in caregiver_prev_shifts:
                prev_shift = caregiver_prev_shifts[caregiver_id]['shift_name']
                if prev_shift in shifts:
                    # 保留在之前的班次，从其他班次移除
                    for shift in shifts:
                        if shift != prev_shift:
                            adjusted_assignments[shift]['caregivers'].remove(caregiver_id)
                    logger.info(f"优化: 保持护理员 {caregiver_id} 在班次 {prev_shift}")
        
        return adjusted_assignments
        
    def _check_consecutive_work_days(self, caregiver_id: str, schedule_date: date) -> bool:
        """
        检查护理员是否连续工作超过6天
        参数：
            caregiver_id: 护理员ID
            schedule_date: 当前排班日期
        返回：
            是否需要休息（True表示需要休息）
        """
        # 查询前6天的排班记录
        start_date = schedule_date - timedelta(days=6)
        
        # 查询这段时间内护理员的排班记录
        records = ShiftAssignment.query.filter(
            ShiftAssignment.caregiver_id == caregiver_id,
            ShiftAssignment.schedule_date >= start_date,
            ShiftAssignment.schedule_date < schedule_date
        ).order_by(ShiftAssignment.schedule_date).all()
        
        # 如果过去6天都有工作，则今天需要休息
        if len(records) >= 6:
            work_dates = set(func.date(record.schedule_date) for record in records)
            consecutive_days = 0
            check_date = schedule_date - timedelta(days=1)
            
            # 检查是否有连续的6天
            while check_date >= start_date:
                if check_date in work_dates:
                    consecutive_days += 1
                    if consecutive_days >= 6:
                        return True  # 需要休息
                else:
                    consecutive_days = 0  # 重置连续工作天数
                check_date -= timedelta(days=1)
        
        return False  # 不需要休息
        
    def _process_consecutive_days(self, start_date: date) -> Dict:
        """
        处理连续多天的排班
        参数：
            start_date: 开始日期
        返回：
            包含所有日期处理结果的字典
        """
        period_dates = self._get_period_dates(start_date)
        results = {}
        
        # 存储每个护理员的工作记录，用于跨天约束处理
        caregiver_work_history = {}
        
        for day_index, schedule_date in enumerate(period_dates):
            try:
                # 检查该日期是否有排班数据
                schedule_exists = ScheduleResult.query.filter(
                    func.date(ScheduleResult.schedule_date) == schedule_date
                ).first() is not None
                
                if not schedule_exists:
                    logger.warning(f"日期 {schedule_date} 没有排班数据，跳过处理")
                    continue
                
                # 1. 为班次分配护理人员
                shift_assignments = self._assign_caregivers_to_shifts(schedule_date)
                
                # 2. 平衡工作量
                shift_assignments = self._balance_workload(shift_assignments)
                
                # 3. 应用白天/黑夜权重
                shift_assignments = self._apply_day_night_weight(shift_assignments)
                
                # 4. 应用CSP约束（连续工作日限制、班次转换规则等）
                shift_assignments = self._apply_csp_constraints(shift_assignments, schedule_date, day_index)
                
                # 5. 预留应急资源
                shift_assignments = self._reserve_emergency_resource(shift_assignments)
                
                # 6. 更新数据库中的时间槽信息
                self._update_time_slots(schedule_date, shift_assignments)
                
                # 记录结果
                results[schedule_date.isoformat()] = {
                    'status': 'success',
                    'shifts': {
                        name: {
                            'caregiver_count': len(data['caregivers']),
                            'reserved_count': len(data.get('reserved_caregivers', set())),
                            'slot_count': len(data['slots'])
                        } for name, data in shift_assignments.items()
                    }
                }
                
                # 更新护理人员工作量统计
                for shift_name, data in shift_assignments.items():
                    for caregiver_id in data['caregivers']:
                        self.caregiver_workload[caregiver_id] = self.caregiver_workload.get(caregiver_id, 0) + 1
                        
                        # 更新工作历史记录
                        if caregiver_id not in caregiver_work_history:
                            caregiver_work_history[caregiver_id] = []
                        caregiver_work_history[caregiver_id].append({
                            'date': schedule_date,
                            'shift': shift_name
                        })
                
                logger.info(f"成功处理日期 {schedule_date} 的排班")
                
            except Exception as e:
                logger.error(f"处理日期 {schedule_date} 的排班时出错: {str(e)}")
                results[schedule_date.isoformat()] = {
                    'status': 'error',
                    'error': str(e)
                }
        
        return results

    def process_schedule(self, start_date=None):
        """
        处理排班计划
        参数：start_date(开始日期，默认为今天)
        返回：处理结果字典
        """
        # 设置默认开始日期
        if start_date is None:
            start_date = date.today()
        elif isinstance(start_date, str):
            start_date = datetime.strptime(start_date, "%Y-%m-%d").date()
        
        # 判断是否为多天排班
        if self.schedule_period_length > 1 or self.schedule_period_type != 'day':
            # 使用连续多天处理方法
            return self._process_consecutive_days(start_date)
        else:
            # 单天处理逻辑
            results = {}
            try:
                # 检查该日期是否有排班数据
                schedule_exists = ScheduleResult.query.filter(
                    func.date(ScheduleResult.schedule_date) == start_date
                ).first() is not None
                
                if not schedule_exists:
                    logger.warning(f"日期 {start_date} 没有排班数据，跳过处理")
                    return {"status": "error", "error": "没有排班数据"}
                
                # 1. 为班次分配护理人员
                shift_assignments = self._assign_caregivers_to_shifts(start_date)
                
                # 2. 平衡工作量
                shift_assignments = self._balance_workload(shift_assignments)
                
                # 3. 应用白天/黑夜权重
                shift_assignments = self._apply_day_night_weight(shift_assignments)
                
                # 4. 预留应急资源
                shift_assignments = self._reserve_emergency_resource(shift_assignments)
                
                # 5. 更新数据库中的时间槽信息
                self._update_time_slots(start_date, shift_assignments)
                
                # 记录结果
                results[start_date.isoformat()] = {
                    'status': 'success',
                    'shifts': {
                        name: {
                            'caregiver_count': len(data['caregivers']),
                            'reserved_count': len(data.get('reserved_caregivers', set())),
                            'slot_count': len(data['slots'])
                        } for name, data in shift_assignments.items()
                    }
                }
                
                logger.info(f"成功处理日期 {start_date} 的排班")
                
            except Exception as e:
                logger.error(f"处理日期 {start_date} 的排班时出错: {str(e)}")
                results[start_date.isoformat()] = {
                    'status': 'error',
                    'error': str(e)
                }
        
        return results 